import streamlit as st
from ultralytics import YOLO
import cv2
from PIL import Image
import os
import numpy as np
from collections import Counter

# --- 系统名称 ---
SYSTEM_NAME = "VisDrone Detector Pro" # 或者你选择的其他名字

# --- Streamlit 应用配置 ---
st.set_page_config(
    page_title=f"{SYSTEM_NAME} - VisDrone Object Detection",
    page_icon="🚁",
    layout="wide"
)

st.title(f"🚁 {SYSTEM_NAME}")
st.markdown(f"""
    **VisDrone Detector Pro** 是一款基于改进的 RT-DETR 模型，用于在 **VisDrone** 数据集上的目标检测可视化工具。
    上传您自己的图像，查看模型如何检测无人机场景中的目标，并了解检测到的目标类别分布。
""")

# --- 模型加载 ---
@st.cache_resource  # 缓存模型，避免重复加载
def load_model(model_path="/home/zhanghao/rtdetr_2/best.pt"):
    """加载 YOLO 模型"""
    if not os.path.exists(model_path):
        st.error(f"模型文件 '{model_path}' 不存在。请确保模型文件已上传或路径正确。")
        return None
    try:
        model = YOLO(model_path)
        st.success("模型加载成功！")
        return model
    except Exception as e:
        st.error(f"加载模型时出错: {e}")
        return None

# --- 图像处理和预测 ---
def predict_objects(image_path, model, conf_threshold=0.3):
    """
    对图像进行目标检测，返回带有边界框的图像和类别统计信息。
    返回: (annotated_image_array, class_counts)
    """
    if model is None:
        return None, {}

    try:
        # 加载图像
        img = Image.open(image_path).convert("RGB")
        img_np = np.array(img)

        # 进行预测
        results = model(img_np, conf=conf_threshold)

        # 处理结果并绘制边界框
        annotated_img = img_np.copy()
        detected_classes = []

        for r in results:
            boxes = r.boxes  # Boxes object for bounding box outputs
            for box in boxes:
                # 提取边界框坐标
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                confidence = float(box.conf[0])
                class_id = int(box.cls[0])
                class_name = model.names[class_id]

                detected_classes.append(class_name) # 记录检测到的类别

                # 绘制边界框和标签
                color = (255, 0, 0) # 蓝色框
                cv2.rectangle(annotated_img, (x1, y1), (x2, y2), color, 2)
                label = f"{class_name}: {confidence:.2f}"
                cv2.putText(annotated_img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

        # 统计类别数量
        class_counts = Counter(detected_classes)

        return annotated_img, class_counts

    except FileNotFoundError:
        st.error(f"图像文件 '{image_path}' 不存在。")
        return None, {}
    except Exception as e:
        st.error(f"预测时出错: {e}")
        return None, {}

# --- Streamlit 界面 ---
if __name__ == "__main__":
    # 加载模型
    model_path = "/home/zhanghao/rtdetr_2/best.pt"  # 确保此路径指向你的 best.pt 文件
    yolo_model = load_model(model_path)

    if yolo_model:
        st.sidebar.header("设置")

        # 可选：上传自定义图像
        uploaded_file = st.sidebar.file_uploader("上传一张 VisDrone 图像进行检测...", type=["jpg", "jpeg", "png"])

        # 可选：选择预设的 VisDrone 图像 (如果你有的话)
        sample_image_dir = "/home/zhanghao/rtdetr_2/sample_images" # 假设你有一个名为 'sample_images' 的文件夹
        sample_images = []
        if os.path.exists(sample_image_dir) and os.path.isdir(sample_image_dir):
            sample_images = [f for f in os.listdir(sample_image_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
            if sample_images:
                selected_sample = st.sidebar.selectbox("或选择一张示例图像:", ["-- 请选择 --"] + sample_images)
            else:
                selected_sample = None
        else:
            selected_sample = None
            st.sidebar.info(f"未找到 '{sample_image_dir}' 文件夹，无法显示示例图像。")


        # 确定要处理的图像路径
        image_to_process = None
        if uploaded_file is not None:
            # 将上传的文件保存到临时目录，以便模型处理
            # 这里的临时文件处理方式可以更 robust，例如使用 tempfile 模块
            image_filename = uploaded_file.name
            with open(image_filename, "wb") as f:
                f.write(uploaded_file.getbuffer())
            image_to_process = image_filename
        elif selected_sample and selected_sample != "-- 请选择 --":
            image_to_process = os.path.join(sample_image_dir, selected_sample)
        else:
            st.warning("请上传一张图像或选择一张示例图像。")

        if image_to_process:
            st.subheader("输入图像")
            st.image(image_to_process, caption="待检测图像", use_column_width=True)

            # 置信度阈值滑块
            confidence_threshold = st.slider("置信度阈值:", 0.0, 1.0, 0.3, 0.05)

            # 执行预测并显示结果
            if st.button("运行检测"):
                with st.spinner("正在进行目标检测..."):
                    annotated_image_array, class_counts = predict_objects(image_to_process, yolo_model, conf_threshold=confidence_threshold)

                if annotated_image_array is not None:
                    st.subheader("检测结果")
                    st.image(annotated_image_array, caption="检测结果", use_column_width=True)

                    # 显示类别统计信息
                    if class_counts:
                        st.subheader("目标类别统计")
                        # 创建一个 Pandas DataFrame 以便更好地显示表格
                        try:
                            import pandas as pd
                            # 对类别按名称排序，或者按数量排序
                            sorted_counts = sorted(class_counts.items(), key=lambda item: item[1], reverse=True) # 按数量降序排序
                            df = pd.DataFrame(sorted_counts, columns=["类别", "数量"])
                            st.dataframe(df)
                        except ImportError:
                            st.warning("请安装 pandas 库 (`pip install pandas`) 以获得更好的表格显示效果。")
                            for class_name, count in class_counts.items():
                                st.write(f"- **{class_name}**: {count}")
                    else:
                        st.info("未检测到任何目标。")
                else:
                    st.error("未能生成检测结果。")

    else:
        st.error("无法加载模型，请检查模型文件路径和完整性。")